#include "compi/compi.h"

#include <iostream>
#include <mpi.h>
#include <string>
#include <vector>

using namespace compi;

static void
run_demo_request_layer(Context & ctx, int rank, TagNamespace ns)
{
    if (rank == 0)
    {
        std::string msg = "req: hello";
        std::string reply;

        SendRequest sreq(ctx, /*dest=*/1, msg, ns);
        RecvRequest rreq(ctx, /*src=*/1, reply, ns);

        std::cout << "[request] rank 0 posted send: " << msg << std::endl;

        ProgressEngine engine(/*max_inflight=*/2);
        engine.add(std::move(sreq));
        engine.add(std::move(rreq));
        while (engine.waitsome() < 2)
            ;

        std::cout << "[request] rank 0 received: " << reply << std::endl;
    }
    else if (rank == 1)
    {
        std::string received;
        RecvRequest rreq(ctx, /*src=*/0, received, ns);
        ProgressEngine engine_recv(/*max_inflight=*/1);
        engine_recv.add(std::move(rreq));
        while (engine_recv.waitsome() < 1)
            ;
        std::cout << "[request] rank 1 received: " << received << std::endl;

        std::string response = received + " | ack";
        SendRequest sreq(ctx, /*dest=*/0, response, ns);
        ProgressEngine engine_send(/*max_inflight=*/1);
        engine_send.add(std::move(sreq));
        while (engine_send.waitsome() < 1)
            ;
        std::cout << "[request] rank 1 sent: " << response << std::endl;
    }
}

static void
run_demo_framed_layer(Context & ctx, int rank, TagNamespace ns)
{
    Endpoint src(ctx.comm(), /*target_rank=*/(rank == 0 ? 1 : 0));
    Endpoint dst(ctx.comm(), /*target_rank=*/(rank == 0 ? 1 : 0));

    FramedChannel framed(/*allow_any_tag=*/true);
    TypedChannel<std::string> tchan(framed);

    if (rank == 0)
    {
        std::string payload = "framed: hi";
        std::string reply;

        auto sres = tchan.send(dst, payload, ns);
        if (sres.is_error())
        {
            std::cerr << "[framed] rank 0 send error: " << sres.error_message() << std::endl;
            return;
        }
        std::cout << "[framed] rank 0 sent: " << payload << std::endl;

        auto rres = tchan.recv(src, reply, ns);
        if (rres.is_error())
        {
            std::cerr << "[framed] rank 0 recv error: " << rres.error_message() << std::endl;
            return;
        }
        std::cout << "[framed] rank 0 received: " << reply << std::endl;
    }
    else if (rank == 1)
    {
        std::string received;
        auto rres = tchan.recv(src, received, ns);
        if (rres.is_error())
        {
            std::cerr << "[framed] rank 1 recv error: " << rres.error_message() << std::endl;
            return;
        }
        std::cout << "[framed] rank 1 received: " << received << std::endl;

        std::string response = received + " | ack";
        auto sres = tchan.send(dst, response, ns);
        if (sres.is_error())
        {
            std::cerr << "[framed] rank 1 send error: " << sres.error_message() << std::endl;
            return;
        }
        std::cout << "[framed] rank 1 sent: " << response << std::endl;
    }
}

int
main()
{
    Environment env;
    env.set_failure_policy(FailurePolicy::ThrowException);
    Context & ctx = Context::for_comm(MPI_COMM_WORLD, env);

    const int rank = ctx.rank();
    const int world_size = ctx.size();
    if (world_size < 2)
    {
        if (rank == 0)
        {
            std::cerr << "This demo requires at least 2 MPI processes." << std::endl;
        }
        return 0;
    }

    // 创建命名空间（用于 ANY_TAG 过滤与标签生命周期管理）
    auto & alloc = ctx.tag_allocator();
    auto ns = alloc.create_namespace("layers_demo");

    // 展示请求层（非阻塞请求/进度引擎）
    run_demo_request_layer(ctx, rank, ns);

    // 展示传输分层（Endpoint + FramedChannel + TypedChannel）
    run_demo_framed_layer(ctx, rank, ns);

    return 0;
}